#include "stdafx.h"
#include "FileSystem.h"

typedef long							NTSTATUS;
#define NT_EXTRACT_ST(Status)			((((ULONG)(Status)) >> 30)& 0x3)
#define NT_SUCCESS(Status)              (((NTSTATUS)(Status)) >= 0)
#define NT_INFORMATION(Status)          (NT_EXTRACT_ST(Status) == 1)
#define NT_WARNING(Status)              (NT_EXTRACT_ST(Status) == 2)
#define NT_ERROR(Status)                (NT_EXTRACT_ST(Status) == 3)

typedef struct _ANSI_STRING {
	USHORT Length;
	USHORT MaximumLength;
	PCHAR Buffer;
} ANSI_STRING, *PANSI_STRING;

typedef struct _OBJECT_ATTRIBUTES {
    HANDLE		 RootDirectory;
	PANSI_STRING ObjectName;
    ULONG		 Attributes;
} OBJECT_ATTRIBUTES, *POBJECT_ATTRIBUTES;

EXTERN_C {
	void RtlInitAnsiString(PANSI_STRING DestinationString, const char* SourceString);
	NTSTATUS ObCreateSymbolicLink(PANSI_STRING, PANSI_STRING);
	NTSTATUS ObDeleteSymbolicLink(PANSI_STRING);
	NTSTATUS NtOpenSymbolicLinkObject(PHANDLE LinkHandle, POBJECT_ATTRIBUTES ObjectAttributes);
	NTSTATUS NtQuerySymbolicLinkObject(HANDLE LinkHandle, PANSI_STRING LinkTarget, PULONG ReturnedLength);
	NTSTATUS NtClose(HANDLE Handle);
}

HRESULT CreateSymbolicLink(const CHAR* szDrive, const CHAR* szDeviceName, BOOL System) {
	// Setup our path
	CHAR szDestinationDrive[MAX_PATH];
	if(System)
		sprintf_s(szDestinationDrive, MAX_PATH, "\\System??\\%s", szDrive);
	else
		sprintf_s(szDestinationDrive, MAX_PATH, "\\??\\%s", szDrive);

	// Setup our strings
	ANSI_STRING linkname, devicename;
	RtlInitAnsiString(&linkname, szDestinationDrive);
	RtlInitAnsiString(&devicename, szDeviceName);

	// Create finally
	NTSTATUS status = ObCreateSymbolicLink(&linkname, &devicename);
	if(NT_SUCCESS(status) || status == 0xc0000035) return S_OK;
	return S_FALSE;
}

HRESULT DeleteSymbolicLink(const CHAR* szDrive, BOOL System) {
	// Setup our path
	CHAR szDestinationDrive[MAX_PATH];
	if(System)
		sprintf_s(szDestinationDrive, MAX_PATH, "\\System??\\%s", szDrive);
	else
		sprintf_s(szDestinationDrive, MAX_PATH, "\\??\\%s", szDrive);

	// Setup our string
	ANSI_STRING linkname;
	RtlInitAnsiString(&linkname, szDestinationDrive);
	
	// Delete finally
	NTSTATUS status = ObDeleteSymbolicLink(&linkname);
	if(NT_SUCCESS(status)) return S_OK;
	return S_FALSE;
}

Drive::Drive(const CHAR* DriveName, const CHAR* PhysicalPath) {
	driveName = DriveName;
	drivePath = driveName + ":";
	physicalPath = PhysicalPath;
	mounted = FALSE;
}
HRESULT Drive::MountForTitle() {
	// Create a symbolic link to this drive
	if(CreateSymbolicLink(drivePath.c_str(), physicalPath.c_str(), FALSE) != S_OK)
		return S_FALSE;

	// Check if it exists now
	string path = drivePath + "\\";
	if(GetFileAttributes(path.c_str()) == 0xFFFFFFFF)
		return S_FALSE;

	// Its mounted so lets set that
	mounted = TRUE;
	return S_OK;
}
HRESULT Drive::Mount() {
	if(CreateSymbolicLink(drivePath.c_str(), physicalPath.c_str(), FALSE) != S_OK)
		return S_FALSE;

	// Check if it exists now
	string path = drivePath + "\\";
	if(GetFileAttributes(path.c_str()) == 0xFFFFFFFF)
		return S_FALSE;

	// Its mounted so lets set that
	mounted = TRUE;
	return S_OK;
}
HRESULT Drive::Unmount() {
	DeleteSymbolicLink(drivePath.c_str(), FALSE);
	mounted = FALSE;
	return S_OK;
}

FileSystem::FileSystem() {
	drives.push_back(new Drive("Sfc",         "\\Device\\Flash"));
	drives.push_back(new Drive("Memunit0",    "\\Device\\Mu0"));
	drives.push_back(new Drive("Memunit1",    "\\Device\\Mu1"));
	drives.push_back(new Drive("OnBoardMU",   "\\Device\\BuiltInMuSfc"));
	drives.push_back(new Drive("DVD",         "\\Device\\Cdrom0"));
	drives.push_back(new Drive("Z",           "\\Device\\Cdrom0"));
	drives.push_back(new Drive("HDD",         "\\Device\\Harddisk0\\Partition1"));
	drives.push_back(new Drive("HDDCache",    "\\Device\\Harddisk0\\Partition0"));
	drives.push_back(new Drive("DEVKIT",      "\\Device\\Harddisk0\\Partition1\\DEVKIT"));
	drives.push_back(new Drive("E",           "\\Device\\Harddisk0\\Partition1\\DEVKIT"));
	drives.push_back(new Drive("Hdd0",        "\\Device\\Harddisk0\\Partition0"));
	drives.push_back(new Drive("HddX",        "\\Device\\Harddisk0\\SystemPartition"));
	drives.push_back(new Drive("USB0",        "\\Device\\Mass0"));
	drives.push_back(new Drive("USB1",        "\\Device\\Mass1"));
	drives.push_back(new Drive("USB2",        "\\Device\\Mass2"));
	drives.push_back(new Drive("USBMU0",      "\\Device\\Mass0PartitionFile\\Storage"));
	drives.push_back(new Drive("USBMU1",      "\\Device\\Mass1PartitionFile\\Storage"));
	drives.push_back(new Drive("USBMU2",      "\\Device\\Mass2PartitionFile\\Storage"));
	drives.push_back(new Drive("USBMUCache0", "\\Device\\Mass0PartitionFile\\StorageSystem"));
	drives.push_back(new Drive("USBMUCache1", "\\Device\\Mass1PartitionFile\\StorageSystem"));
	drives.push_back(new Drive("USBMUCache2", "\\Device\\Mass2PartitionFile\\StorageSystem"));
	drives.push_back(new Drive("SAPTest",     "\\Device\\Harddisk0\\SystemAuxPartition"));
	drives.push_back(new Drive("Game",		  "??"));
	drives.push_back(new Drive("D",		      "??"));
}
HRESULT FileSystem::MountDrives() {
	
	// Loop through all our drives and mount them
	for(DWORD x = 0; x < drives.size(); x++)
		drives[x]->Mount();

	// All done
	return S_OK;
}
HRESULT FileSystem::UnMountDrives() {
	for(DWORD i = 0;i < drives.size();i++)
		drives[i]->Unmount();

	return S_OK;
}
HRESULT FileSystem::VerifyDrives() {
	// Check drives, mount any that aren't mounted
	for(DWORD i = 0;i < drives.size();i++) {
		/*if(!*/drives[i]->IsMounted();/*)*/
			//drives[i]->Mount();
	}
	return S_OK;
}
HRESULT FileSystem::MountForTitle(const char* Path) {

	// Get our path as a string
	string path = Path;
	path = path.substr(0, path.find_first_of(":"));

	// Now figure out the drive it is so we can mount it
	for(DWORD x = 0; x < drives.size(); x++)
		if(stricmp(path.c_str(), drives[x]->GetDriveName()) == 0)
			return drives[x]->MountForTitle();

	// Hmmm we didnt find it..
	return S_FALSE;
}